# agents/planner_agent.py
# -*- coding: utf-8 -*-
from autogen_agentchat.agents import AssistantAgent
from autogen_core import CancellationToken


class PlannerAgent(AssistantAgent):
    """
    PlannerAgent：负责将用户自然语言需求拆解成 MCP 工具调用 + SQL。
    通过注入的 MCP 工具列表（在 AssistantAgent 构造时传入 tools）来调用 get_financial_tables/execute_sql。
    """

    def __init__(self, *args, **kwargs):
        # 保留继承构造（从外部传入 model_client 与 tools）
        super().__init__(*args, **kwargs)
        # 你可以在 system_message 中约束 LLM 的输出格式，便于后续解析
        self._system_message = (
            "你是调度助手。用户会发送类似：'帮我分析 2024年1月到6月的财务情况' 的需求。"
            "请按步骤：1) 调用工具 get_financial_tables(keyword=...) 询问相关表。"
            "2) 根据返回的表结构生成一条明确的 SQL 语句（只输出 SQL，不要额外解释）。"
            "3) 指示使用工具 execute_sql(sql=...) 执行。"
            "注意：工具调用格式必须严格为 JSON 字典，便于上层自动调用。"
        )

    async def on_messages(self, messages, cancellation_token: CancellationToken):
        """
        我们可以复写 on_messages 来体现一个明确流程：先询问表，再生成 SQL，再执行查询。
        但因为 AssistantAgent 已有工具集成能力（会自动把模型里生成的 function/tool call 转换为工具调用）
        这里仅给出一个最小的交互流示例（简化流程，部分由模型生成指令）。
        """
        # 简单把最后一条用户消息取出并交给基类处理（基类会使用 model_client + tools）
        return await super().on_messages(messages, cancellation_token)
